{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nested Context Managers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the last video we saw that we could nest context managers.\n",
    "\n",
    "This is actually fairly common.\n",
    "\n",
    "Suppose we need to open a number of files - using a `with` statement for each one means we would have to nest that many `with` statements as well.\n",
    "\n",
    "For example, we want to \"zip\" three files. Let's look at the content of each file first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file1_line1\n",
      "file1_line2\n",
      "file1_line3\n",
      "----------------\n",
      "file2_line1\n",
      "file2_line2\n",
      "file2_line3\n",
      "----------------\n",
      "file3_line1\n",
      "file3_line2\n",
      "file3_line3"
     ]
    }
   ],
   "source": [
    "with open('file1.txt') as f:\n",
    "    for row in f:\n",
    "        print(row, end='')\n",
    "print('\\n----------------')\n",
    "with open('file2.txt') as f:\n",
    "    for row in f:\n",
    "        print(row, end='')\n",
    "print('\\n----------------')\n",
    "with open('file3.txt') as f:\n",
    "    for row in f:\n",
    "        print(row, end='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to combine the rows from each file, and print them out - like zipping together basically, except we need to strip out that pesky `\\n`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n"
     ]
    }
   ],
   "source": [
    "with open('file1.txt') as f1:\n",
    "    with open('file2.txt') as f2:\n",
    "        with open('file3.txt') as f3:\n",
    "            while True:\n",
    "                try:\n",
    "                    f1_row = next(f1).strip('\\n')\n",
    "                    f2_row = next(f2).strip('\\n')\n",
    "                    f3_row = next(f3).strip('\\n')\n",
    "                except StopIteration:\n",
    "                    break\n",
    "                else:\n",
    "                    print(f1_row + ',' + f2_row + ',' + f3_row)\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see we needed three levels of nested `with` statements.\n",
    "\n",
    "Instead, we might try to approach the problem this way - but first let's write our own `openfile` context manager so we can easily see when the file is being opened and closed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "\n",
    "@contextmanager\n",
    "def open_file(f_name):\n",
    "    print(f'opening file {f_name}')\n",
    "    f = open(f_name)\n",
    "    try:\n",
    "        yield f\n",
    "    finally:\n",
    "        print(f'closing file {f_name}')\n",
    "        f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we are going to create (but not enter) the context managers, and store the enter and exit methods in some lists. We'll also create a list that will contain the values returned by the enter methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "enters = []\n",
    "exits = []\n",
    "for f_name in f_names:\n",
    "    ctx = open_file(f_name)\n",
    "    enters.append(ctx.__enter__)\n",
    "    exits.append(ctx.__exit__)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we are going to enter the contexts by calling the `__enter__` methods, store the return values (the file objects), process the data, and then run all the `__exit__` methods (in reverse order!):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n"
     ]
    }
   ],
   "source": [
    "files = [enter() for enter in enters]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    try:\n",
    "        rows = [next(f).strip('\\n') for f in files]\n",
    "    except StopIteration:\n",
    "        break\n",
    "    else:\n",
    "        row = ','.join(rows)\n",
    "        print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to close all the files by calling the `__exit__` methods (in reverse order, since we aded them in the order in which the contexts were opened (i.e. we open from first file to last, but close from last opened file to first - think of how the context managers are nested).\n",
    "\n",
    "Also, keep in mind that `__exit__` methods need those exception parameters - here we'll just use None for simplicity - we are not doing any exception handling!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "for fn in exits[::-1]:\n",
    "    fn(None, None, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, let's recap by putting all this code together and simplifying a few things along the way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "exits = []\n",
    "files = []\n",
    "try:\n",
    "    for f_name in f_names:\n",
    "        ctx = open_file(f_name)\n",
    "        files.append(ctx.__enter__())\n",
    "        exits.append(ctx.__exit__)    \n",
    "    \n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)\n",
    "finally:\n",
    "    for fn in exits[::-1]:\n",
    "        fn(None, None, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What was simpler, this method or simply nesting the `with` blocks?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if we were doing this with 100 files instead of just 3?\n",
    "\n",
    "Or what if we did not know in advance how many files we had to zip together (maybe we're reading all the .txt files in a directory for example - there may be 1 file, 3 files, 10 files, we don't realy know, and it can change over time)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Maybe we can find a way to make this a little easier to use.\n",
    "\n",
    "Let's try using a context manager to hold on to all these `__exit__` methods we want to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NestedContexts:\n",
    "    def __init__(self, *contexts):\n",
    "        self._enters = []\n",
    "        self._exits = []\n",
    "        self._values = []\n",
    "        \n",
    "        for ctx in contexts:\n",
    "            self._enters.append(ctx.__enter__)\n",
    "            self._exits.append(ctx.__exit__)\n",
    "        \n",
    "    def __enter__(self):\n",
    "        for enter in self._enters:\n",
    "            self._values.append(enter())\n",
    "        return self._values\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, exc_tb):\n",
    "        for exit in self._exits[::-1]:\n",
    "            exit(exc_type, exc_value, exc_tb)\n",
    "        return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try to use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "do work here\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "with NestedContexts(open_file('file1.txt'),\n",
    "                   open_file('file2.txt'),\n",
    "                   open_file('file3.txt')) as files:\n",
    "    print('do work here')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the files were opened, our work was done, and the files were closed.\n",
    "Let's just add in some real work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "with NestedContexts(open_file('file1.txt'),\n",
    "                   open_file('file2.txt'),\n",
    "                   open_file('file3.txt')) as files:\n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is much better, but specifying the context managers is still a little painful, having to list them all separately as the arguments to the `NestedContexts` manager.\n",
    "\n",
    "We could simplify things somewhat by taking this approach:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "file_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "contexts = [open_file(f_name) for f_name in f_names]\n",
    "with NestedContexts(*contexts) as files:\n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, this works, and is actually quite workable, but we have to do some setup work before we can use the context manager.\n",
    "\n",
    "We can try a slightly different approach where we create a method in our `NestedContextManager` that can be used to \"register\" contexts - so instead of creating the `NestedContextManager` with an `__init__` that takes in all the contexts at once, we create the `NextedContextManager` **first**, and then, inside the `with` block we append the contexts we want to work with.\n",
    "\n",
    "One main advantage to that approach is that we can add contexts to `NestedContextManager` at any time in the `with` block - i.e. we can delay when and how we add contexts to the overarching context manager.\n",
    "\n",
    "So, we'll do this by implementing a method in `NestedContexts` itself that will allow us to append a context manager, get the `__enter__` value out of it, and store the `__exit__` methods.\n",
    "\n",
    "To do this we're going to take a slightly different approach - the `NestedContexts` manager is going to return itself in it's `__enter__` method, instead of returning a list of the various context values returned from their respective `__enter__` methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NestedContexts:\n",
    "    def __init__(self):\n",
    "        self._exits = []\n",
    "        \n",
    "    def __enter__(self):\n",
    "        return self\n",
    "\n",
    "    def enter_context(self, ctx):\n",
    "        self._exits.append(ctx.__exit__)\n",
    "        value = ctx.__enter__()\n",
    "        return value\n",
    "        \n",
    "    def __exit__(self, exc_type, exc_value, exc_tb):\n",
    "        for exit in self._exits[::-1]:\n",
    "            exit(exc_type, exc_value, exc_tb)\n",
    "        return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's try it again, but this time we'll \"register\" our contexts once the `NestedContexts` context has been entered:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "with NestedContexts() as stack:\n",
    "    files = [stack.enter_context(open_file(f_name)) for f_name in f_names]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! Now let's just do the work as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "with NestedContexts() as stack:\n",
    "    files = [stack.enter_context(open_file(f_name)) for f_name in f_names]\n",
    "    \n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully you can now see why I said we can decide to add contyexts to that stack at any time inside the `with` statement - we are not restricted to adding them in the `__init__` - which means we can even use `if` statements to add contexts to the stack if we want to - this is far more flexible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, this is a common enough scenario that the standard library has something up its sleeve for us!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `contextlib` has an `ExitStack` context manager that works the same way as our `NestedContexts`, but, unlike our approach, it actually does exception handling properly too!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how we use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from contextlib import ExitStack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "with ExitStack() as stack:\n",
    "    files = [stack.enter_context(open_file(f_name))\n",
    "            for f_name in f_names]    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the files were opened and automatically closed. Now all we need to do is the work itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file file1.txt\n",
      "opening file file2.txt\n",
      "opening file file3.txt\n",
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n",
      "closing file file3.txt\n",
      "closing file file2.txt\n",
      "closing file file1.txt\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "with ExitStack() as stack:\n",
    "    files = [stack.enter_context(open_file(f_name))\n",
    "            for f_name in f_names]\n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To finish up we can use the built-in `open` context manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file1_line1,file2_line1,file3_line1\n",
      "file1_line2,file2_line2,file3_line2\n",
      "file1_line3,file2_line3,file3_line3\n"
     ]
    }
   ],
   "source": [
    "f_names = 'file1.txt', 'file2.txt', 'file3.txt'\n",
    "\n",
    "with ExitStack() as stack:\n",
    "    files = [stack.enter_context(open(f_name))\n",
    "            for f_name in f_names]\n",
    "    while True:\n",
    "        try:\n",
    "            rows = [next(f).strip('\\n') for f in files]\n",
    "        except StopIteration:\n",
    "            break\n",
    "        else:\n",
    "            row = ','.join(rows)\n",
    "            print(row)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
